from typing import Union
from fastapi import APIRouter, HTTPException, BackgroundTasks
from app.services.dhged_api_service import DHGEDApiService
from app.schemas.dhgam.dhged_schema import (
    DHGEDFitRequest, DHGEDDetectRequest, DHGEDExplainRequest,
    DHGEDResponse, DHGEDMiddleResponse
)
from app.schemas.facenet_base_schema import TaskStatus
from app.tools.redis_client import redis_client
import logging

# 配置日志
logger = logging.getLogger(__name__)
router = APIRouter()
service = DHGEDApiService()

@router.post("/create", response_model=DHGEDMiddleResponse)
async def create_fit_task(request: DHGEDFitRequest) -> DHGEDMiddleResponse:
    """
    创建模型训练任务
    
    Args:
        request: 包含图序列和模型参数的请求
        
    Returns:
        DHGEDMiddleResponse: 包含任务初始状态的响应
    """
    try:
        task_response = DHGEDMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params,
            metrics=[]
        )
        await redis_client.set_data(request.task_id, task_response.dict())
        return task_response
    except Exception as e:
        logger.error(f"创建训练任务失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"创建训练任务失败: {str(e)}")

@router.post("/create_detect_task", response_model=DHGEDMiddleResponse)
async def create_detect_task(request: DHGEDDetectRequest) -> DHGEDMiddleResponse:
    """
    创建事件检测任务
    
    Args:
        request: 包含新图和历史图的请求
        
    Returns:
        DHGEDMiddleResponse: 包含任务初始状态的响应
    """
    try:
        task_response = DHGEDMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params,
            metrics=[]
        )
        await redis_client.set_data(request.task_id, task_response.dict())
        return task_response
    except Exception as e:
        logger.error(f"创建检测任务失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"创建检测任务失败: {str(e)}")

@router.post("/create", response_model=DHGEDMiddleResponse)
async def create_explain_task(request: DHGEDExplainRequest) -> DHGEDMiddleResponse:
    """
    创建事件解释任务
    
    Args:
        request: 包含事件的请求
        
    Returns:
        DHGEDMiddleResponse: 包含任务初始状态的响应
    """
    try:
        task_response = DHGEDMiddleResponse(
            task_id=request.task_id,
            task_callback_url=request.task_callback_url,
            task_status=TaskStatus.PENDING,
            task_progress=0,
            input_params=request.input_params,
            metrics=[]
        )
        await redis_client.set_data(request.task_id, task_response.dict())
        return task_response
    except Exception as e:
        logger.error(f"创建解释任务失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"创建解释任务失败: {str(e)}")

async def add_task_to_background(
    background_tasks: BackgroundTasks,
    task_type: str,
    input_params: dict,
    task_id: str
) -> None:
    """将任务添加到后台执行队列"""
    if task_type == "fit":
        background_tasks.add_task(run_fit_task, input_params, task_id)
    elif task_type == "detect":
        background_tasks.add_task(run_detect_task, input_params, task_id)
    elif task_type == "explain":
        background_tasks.add_task(run_explain_task, input_params, task_id)

@router.post("/execute_fit/{task_id}", response_model=DHGEDMiddleResponse)
async def execute_fit_task(
    task_id: str, 
    background_tasks: BackgroundTasks
) -> DHGEDMiddleResponse:
    """
    执行模型训练任务
    
    Args:
        task_id: 任务ID
        background_tasks: 后台任务
        
    Returns:
        DHGEDMiddleResponse: 任务执行状态
    """
    return await execute_task(task_id, background_tasks, "fit")

@router.post("/execute_detect/{task_id}", response_model=DHGEDMiddleResponse)
async def execute_detect_task(
    task_id: str, 
    background_tasks: BackgroundTasks
) -> DHGEDMiddleResponse:
    """
    执行事件检测任务
    
    Args:
        task_id: 任务ID
        background_tasks: 后台任务
        
    Returns:
        DHGEDMiddleResponse: 任务执行状态
    """
    return await execute_task(task_id, background_tasks, "detect")

@router.post("/execute_explain/{task_id}", response_model=DHGEDMiddleResponse)
async def execute_explain_task(
    task_id: str, 
    background_tasks: BackgroundTasks
) -> DHGEDMiddleResponse:
    """
    执行事件解释任务
    
    Args:
        task_id: 任务ID
        background_tasks: 后台任务
        
    Returns:
        DHGEDMiddleResponse: 任务执行状态
    """
    return await execute_task(task_id, background_tasks, "explain")

async def execute_task(
    task_id: str, 
    background_tasks: BackgroundTasks,
    task_type: str
) -> DHGEDMiddleResponse:
    """执行指定类型的任务"""
    try:
        task_data = await redis_client.get_data(task_id)
        if not task_data:
            raise HTTPException(status_code=404, detail="任务不存在")

        input_params = task_data.get("input_params")
        result_key = f'{task_id}_result'
        
        # 更新任务状态为执行中
        running_status = {
            "task_id": task_id,
            "task_callback_url": task_data.get("task_callback_url"),
            "task_status": TaskStatus.RUNNING,
            "task_progress": 0,
            "input_params": input_params,
            "metrics": []
        }
        await redis_client.set_data(result_key, running_status)
        
        # 添加到后台任务队列
        await add_task_to_background(background_tasks, task_type, input_params, task_id)
        
        return DHGEDMiddleResponse(**running_status)
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"执行任务失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"执行任务失败: {str(e)}")

@router.get("/get_result/{task_id}", response_model=Union[DHGEDResponse, DHGEDMiddleResponse])
async def get_result(task_id: str) -> Union[DHGEDResponse, DHGEDMiddleResponse]:
    """
    获取任务执行结果
    
    Args:
        task_id: 任务ID
        
    Returns:
        Union[DHGEDResponse, DHGEDMiddleResponse]: 任务结果或状态
    """
    try:
        result_key = f'{task_id}_result'
        result_data = await redis_client.get_data(result_key)
        if not result_data:
            raise HTTPException(status_code=404, detail="结果不存在")
        
        return (DHGEDResponse(**result_data) 
                if "output_params" in result_data or "error_message" in result_data
                else DHGEDMiddleResponse(**result_data))
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取结果失败: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"获取结果失败: {str(e)}")

async def run_fit_task(input_params: dict, task_id: str) -> None:
    """在后台执行模型训练任务"""
    try:
        result = await service.fit_model(input_params, task_id)
        await redis_client.update_data(f'{task_id}_result', result)
    except Exception as e:
        logger.error(f"训练任务执行失败: {str(e)}", exc_info=True)
        error_status = {
            "task_id": task_id,
            "task_status": TaskStatus.FAILED,
            "error_message": str(e)
        }
        await redis_client.update_data(f'{task_id}_result', error_status)

async def run_detect_task(input_params: dict, task_id: str) -> None:
    """在后台执行事件检测任务"""
    try:
        result = await service.detect_events(input_params, task_id)
        await redis_client.update_data(f'{task_id}_result', result)
    except Exception as e:
        logger.error(f"检测任务执行失败: {str(e)}", exc_info=True)
        error_status = {
            "task_id": task_id,
            "task_status": TaskStatus.FAILED,
            "error_message": str(e)
        }
        await redis_client.update_data(f'{task_id}_result', error_status)

async def run_explain_task(input_params: dict, task_id: str) -> None:
    """在后台执行事件解释任务"""
    try:
        result = await service.explain_event(input_params, task_id)
        await redis_client.update_data(f'{task_id}_result', result)
    except Exception as e:
        logger.error(f"解释任务执行失败: {str(e)}", exc_info=True)
        error_status = {
            "task_id": task_id,
            "task_status": TaskStatus.FAILED,
            "error_message": str(e)
        }
        await redis_client.update_data(f'{task_id}_result', error_status) 